import numpy as np
import cv2,h5py
import os
from getStrokeLabel import genStrokes
from config import Config
config = Config()
draw = genStrokes(config, )


#image_name = os.getcwd() + '/image.png'
#alpha_name = os.getcwd() + '/alpha.png'
#image = cv2.imread(image_name)
#alpha = cv2.imread(alpha_name)
#image=cv2.resize(image,(400,400))
#alpha=cv2.resize(alpha,(400,400))
#alpha = alpha[:,:,0] 
#sequence_length=2


regions_list = []
for x in range(-10, 10, int(20*config.win_img)):
    for y in range(-10, 10, int(20*config.win_img)):
        regions_list.append([x*1.0/10,y*1.0/10])
regions = np.array(regions_list)
######################################################################
def matteSolver(time_step,counter):
    im_sz = config.img_sz
    im_ch = config.img_channels
    output_ch = config.alphaMatte_channels
    cnt = counter
    strokemapname = './strokemaps/stroke_' + str(cnt) +'.png'
    systemCall = './SharedMatting -i ./image.png -t ./strokemaps/stroke_' + str(cnt) +'.png -a'
    alphaname = './results/' + 'matte_' +str(time_step)+'_' + str(cnt) + '.png'
    while(not os.path.exists(strokemapname)):
        print ('matteSolver: %d'%cnt)
        time.sleep(0.1)
    os.system(systemCall)
    matte = cv2.imread('shared_matting_alpha.bmp')
    cv2.imwrite(alphaname, matte)
    matte = cv2.imread(alphaname, 0)
    matte = matte.reshape((1, im_sz, im_sz, 1))
    return matte
######################################################################
def min_max_regions(time_step,selected_strokemap=None):
        mse_1_arr = []
        cnt = 0
        window_num=1/config.win_img
        while(cnt < window_num*window_num):
            loc = regions[cnt]
            cnt += 1
            l_1_j = loc
            strokemap_1_j = draw.drawstroke(selected_strokemap.copy(), l_1_j, alpha, cnt)
            strokemap_1_j = np.uint8(strokemap_1_j)
            strokename = './strokemaps/' + 'stroke_'+ str(cnt) + '.png'
            cv2.imwrite(strokename, strokemap_1_j)
        cnt1 = 0
        while(cnt1 < window_num*window_num):
            cnt1 += 1
            matteSolver(time_step,cnt1)
        cnt2 = 0
        while(cnt2 < window_num*window_num):
            cnt2 += 1
            alphaName = './results/' + 'matte_' +str(time_step)+'_' + str(cnt2) + '.png'
            alphaPNG = cv2.imread(alphaName)
            alpha_1_j = alphaPNG[:,:,0]
            alpha_1_j = alpha_1_j.astype(np.float)
            alpha_1_j = np.expand_dims(alpha_1_j, 0)
            mse_1_j = np.sqrt(np.sum(np.square(alpha - alpha_1_j)))
            mse_1_arr.append(mse_1_j)
        min_index = np.argmin(mse_1_arr)
        min_mse_1 = mse_1_arr[min_index]
        return min_mse_1, min_index
######################################################################
def optimal_seq(index_arr):
    max_length = sequence_length
    idx = 0
    while(idx < max_length):
        idx += 1
        if idx == 1:
            min_mse_1, min_index = min_max_regions(time_step=idx,selected_strokemap=image)
            selected_index = min_index
            print('selected index: %d' % selected_index)
            index_arr.append(selected_index)
            selected_strokename = './strokemaps/' + 'stroke_' + str(selected_index+1) + '.png'
            selected_strokemap = cv2.imread(selected_strokename)

        else:
            min_mse_i, min_index = min_max_regions(time_step=idx,selected_strokemap=selected_strokemap)
            selected_index = min_index
            index_arr.append(selected_index)
            print('selected index: %d' % selected_index)
            selected_strokename = './strokemaps/' + 'stroke_' + str(selected_index+1) + '.png'
            selected_strokemap = cv2.imread(selected_strokename)


#################################################################################
def transform(regions,index):
    return regions[index]



if __name__ == '__main__':
    #image_name = os.getcwd() + '/image.png'
    #alpha_name = os.getcwd() + '/alpha.png'
    #image = cv2.imread(image_name)
    #alpha = cv2.imread(alpha_name)
    #image=cv2.resize(image,(400,400))
    #alpha=cv2.resize(alpha,(400,400))
    #alpha = alpha[:,:,0] 
    
    images=np.array(h5py.File('./train.h5','r').get('images'))
    alphas=np.array(h5py.File('./train.h5','r').get('labels'))
    sequence_length=config.num_glimpses
    index_feature=[]
    if (os.path.exists('./locs.h5')):
        point_index = 99 
        with h5py.File('./locs.h5') as hf:
             index_feature = hf[str(point_index)].value
    

#############################################################################
    #for (image, alpha) in zip(images,alphas):
    #     cv2.imwrite('image.png',image)
    #     index_arr=[]
    #     optimal_seq(index_arr)
    #     print ('result is :', index_arr)
    #     index_feature.append(index_arr)
    #     print ('index feature shape: ',np.shape(index_feature))

    #with h5py.File('./train.h5', 'r+') as hf:
    #      hf.create_dataset('index', data=np.array(index_feature))
    #indexes=np.array(index_feature)
    #locs=np.zeros((indexes.shape[0],indexes.shape[1],2))

    #for i in range(indexes.shape[0]):
    #    for j in range (indexes.shape[1]):
    #       locs[i,j]=transform(regions,indexes[i,j])
    #with h5py.File('./train.h5', 'r+') as hf:
    #      hf.create_dataset('locs', data=np.array(locs))
###############################################################################

    for i in range(point_index+1,images.shape[0]):
        image = images[i]
        alpha = alphas[i]
        cv2.imwrite('image.png',image)
        index_arr=[]
        optimal_seq(index_arr)
        print ('result is :', index_arr)
        index_feature.append(index_arr)
        print ('index feature shape: ',np.shape(index_feature))
        with h5py.File('./locs.h5', 'a') as hf:
          hf[str(i)]=np.array(index_feature)
